package io.yyxx.boot.framework.aop;

import com.google.common.util.concurrent.RateLimiter;
import io.yyxx.boot.framework.annotation.RateLimit;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

@Slf4j
@Aspect
@Component
public class RateLimitAspect {

  private static final ConcurrentMap<String, RateLimiter> rateLimiterCache = new ConcurrentHashMap<>();

  @Around("@within(rateLimit) || @annotation(rateLimit)")
  public Object rateLimit(ProceedingJoinPoint point, RateLimit rateLimit) throws Throwable {
    MethodSignature signature = (MethodSignature) point.getSignature();
    Method method = signature.getMethod();
    String key = method.getDeclaringClass().getName() + "." + method.getName();

    // 获取或创建对应的 RateLimiter 对象
    RateLimiter rateLimiter = rateLimiterCache.computeIfAbsent(key,
        k -> RateLimiter.create(rateLimit.permitsPerSecond()));
    if (rateLimit.timeout() > 0) {
      // 尝试在指定的超时时间内获取令牌
      if (!rateLimiter.tryAcquire(rateLimit.timeout(), rateLimit.timeUnit())) {
        log.warn("请求被限流，错误信息: {}", rateLimit.msg());
        throw new RuntimeException(rateLimit.msg());
      }
    } else {
      // 直接获取令牌，如果没有获取到则阻塞等待
      rateLimiter.acquire();
    }
    return point.proceed();
  }
}
